/*---------------------------------------------------------------------
基于ASR的演示系统demo应用
soc_asr:soc result send by asr
Writed by mtd at 2025.4.3
Modify at 2025.4.10
-----------------------------------------------------------------------*/
#include <iostream>
#include "client_t.h"
#include <cstdlib> //包含system函数
#include <unistd.h>
#include <signal.h>
#include <vector>
#include "taiic_key.h"
#include "taiic_sensor.h"
#include <thread>
#include "taiic_avsr.h"
#include "taiic_mfcc_v3.h"
#include "../../include/rs_485/485_common.h" //485协议头文件
#include "../../include/isp/taiic_isp_model.h"
#include "../../include/taiic_uart/taiic_uart_sum.h"
#include "../../include/mfccs/taiic_mfcc_v3.h"
#include <Eigen/Dense>

using namespace std;
RK_BOOL is_model_run = RK_FALSE;
long pcm_size = 0;

FILE *asr_fp_result = NULL;
char asrCmdFinal_result[256] = {0};


TAIIC_MODE_STATE_CTX_S *mode_state; // 模式状态相关参数
AVSR_TOOLKIT_MODEL_CTX_S *avsr_ctx; // avsr模型推理相关

RK_U8 avsr_audio_input_data[ASR_BATCH_N][ASR_FRAMES_H][ASR_MFCC_W][ASR_CHANNEL_C] = {0}; // 1*48*64*1
static void app_key_callback(TAIIC_KEY_STATE old_state, TAIIC_KEY_NAME key_name, TAIIC_KEY_STATE key_state);
void *ai_uart_thread_func(void *args); // 副板->主板气导音频数据采集
static void aivi_save_end(int vi, int ai); // asr
/*---------------------------------------------------------------------
                            全局变量/宏定义/结构体/函数声明
-----------------------------------------------------------------------*/
static void sigterm_handler(int sig);
// int cnt = 0;
/*----------------------------------------------------------------
                            主函数main
-----------------------------------------------------------------*/
int main(int argc, char **argv)
{
    //UART串口初始化
    serial_init();
    //initial mode state
    mode_state = reinterpret_cast<TAIIC_MODE_STATE_CTX_S *>(malloc(sizeof(TAIIC_MODE_STATE_CTX_S)));
    memset(mode_state, 0, sizeof(TAIIC_MODE_STATE_CTX_S));
    //initial avsr model
    avsr_ctx = reinterpret_cast<AVSR_TOOLKIT_MODEL_CTX_S *>(malloc(sizeof(AVSR_TOOLKIT_MODEL_CTX_S))); // 分配内存空间
    memset(avsr_ctx, 0, sizeof(AVSR_TOOLKIT_MODEL_CTX_S));
    //config avsr model
    avsr_ctx->modelPath = TAIIC_ASR_MODEL_PATH_505;
    // avsr_ctx->modelPath = argv[2];

    avsr_rknn_toolkit_config_init(avsr_ctx);
    avsr_rknn_toolkit_io_init(avsr_ctx);
    RK_LOGD("====taiic init model config successfully===\n");

    //按键检测
    taiic_all_key_registers(mode_state, app_key_callback);
    /*---------------------soc connect---------------------*/
    char *serverAddr = argv[1]; // 服务器的ip地址
    printf("set server add is %s\n", serverAddr);
    // 初始化socket、连接服务器
    sock_cli = init_socket(serverAddr);
    
    DIR *myavidir = NULL;
    if((myavidir = opendir(aviCmd)) == NULL) //创建文件夹avi
    {
        bool ret = mkdir(aviCmd, 0755);
        RK_LOGI("%s created sucess!/n", aviCmd);
    }
    //ai_save end_callback
    aivi_save_end_callback(aivi_save_end);
    //audio_capture
    pthread_t tid_ai_uart;
    pthread_create(&tid_ai_uart, NULL, ai_uart_thread_func, NULL);

    //spp receive data from host
    uart_main_func(10, 0, 0, 0);
    pthread_join(tid_ai_uart, NULL);

    // setup the SIGINT to ctrl+c and handing de-init process
    signal(SIGINT, sigterm_handler);
__FAILED_VSR:
    if (avsr_ctx)
    {
        free(avsr_ctx);
        avsr_ctx = RK_NULL;
    }
__FAILED_SYS:
    RK_MPI_SYS_Exit();
    if (mode_state)
    {
        free(mode_state);
        mode_state = RK_NULL;
    }
    // 释放循环缓冲区
    serialData.releaseRingbuffer();
    return 0;
}
/*---------------------------------------------------------------
                          函数实现
-----------------------------------------------------------------*/
static void sigterm_handler(int sig)
{
    RK_LOGI("Catched SIGINT %d\n", sig);
}


static void app_key_callback(TAIIC_KEY_STATE old_state, TAIIC_KEY_NAME key_name, TAIIC_KEY_STATE key_state)
{
    // 判断key_name去实现对应功能,key_name=0表示音频线模式
    // RK_LOGI("key_name, key_state,count is [%d, %d,%d]\n", key_name, key_state, cnt++);

    switch (key_name)
    {
    case TAIIC_PTT:
        mode_state->ptt_flag = key_state;
        if(mode_state->ptt_flag == PTT_DOWN)
        {
            //特殊模式
            if(!is_model_run)
            {
                //fp
                snprintf(asrCmdFinal_result, sizeof(asrCmdFinal_result), "%ssession%d_%d.txt", aviCmd, sessionNum, typeNUm);
                printf("=====vi_result cmd is %s===\n", asrCmdFinal_result);

                asr_fp_result = fopen(asrCmdFinal_result, "wb"); 
                //PTT按下,触发音频采集
                mtx_uart.lock();
                uart_main_func(1, 1, 0, 0); // SPK嘀一声,开始采集
                mtx_uart.unlock();

                {
                    std::lock_guard<std::mutex> lock(mtx_ai);
                    ai_start = true; // 设置标志位为true，允许ai开始执行
                }
                cv_ai.notify_one(); // 通知等待的线程，条件已经满足

                typeNUm++;
                if (typeNUm == TYPE)
                {
                    sessionNum++;
                    typeNUm = 0;
                }
                if (sessionNum == SESSION)
                {
                    sessionNum = 0;
                    exit(0);
                }
            }
        }
        break;
    }
}

void *ai_uart_thread_func(void *args) //副板->主板气导音频数据采集
{
    while (true)
    {
        std::unique_lock<std::mutex> lock(mtx_ai);
        // 等待ai_start_flag变为true
        cv_ai.wait(lock, []()
                   { return ai_start; });
        
        //清除工作可用的标志，为下一次循环做准备
        ai_start = false;
        // 开始执行任务
        if (fileExists(air_pcm_path))
        {
            if (remove(air_pcm_path) == 0)
            {
                printf("air pcm file exists, and delete successfully\n");
            }
            else
            {
                printf("air pcm file delete failed\n");
            }
        }

        //开始执行任务
        mtx_uart.lock();
        uart_main_func(5, 1, 0, 0);
        mtx_uart.unlock();

        while(1)
        {
            pcm_size = getFileSize(air_pcm_path);
            if(pcm_size == PCM_SIZE_SOC) // 53280 bytes
            {
                ai_end = 1; // ai finish flag
                RK_LOGD("=====save ai finish==\n");

                //end save ai -> go to mfcc,ai_data in /data/air_audio_data.pcm 16k/16bits/single channel
                aivi_save_end(vi_end, ai_end);
                break;
            }
            sleep(1);
        }
    }
    pthread_exit(NULL);
}

static void aivi_save_end(int vi, int ai) // asr
{
    // load audio pcm as input
    audioData = loadAudioData(air_pcm_path);
    // mfcc feature extraction, output type:uint8
    extract_mfcc(audioData, mfcc_feat_u8);
    // 将matrix中的数据赋值到avsr_audio_input_data
    for (int i = 0; i < numFrames; i++)
    {
        for (int j = 0; j < num_mfcc_coef; j++)
        {
            avsr_audio_input_data[0][i][j][0] = static_cast<RK_U8>(mfcc_feat_u8(i, j));
        }
    }

    // 送入模型
    unsigned char *avsr_input_data[avsr_ctx->io_num.n_input];

    avsr_input_data[0] = new unsigned char[avsr_ctx->input_attrs[0].size]; // 设置audio 输入大小
    avsr_input_data[0] = &avsr_audio_input_data[0][0][0][0];

    RK_LOGI("===MFCC END==\n");
    avsr_rknn_toolkit_data_refresh(avsr_ctx, avsr_input_data);
    int ret = rknn_run(avsr_ctx->context, NULL);
    MODEL_RESULT_S avsr_result = avsr_rknn_toolkit_result_int8(avsr_ctx);

    printf("===result label is %d, prob is %f====\n", avsr_result.label, avsr_result.prob);

    fprintf(asr_fp_result, "%03d,%.4f\n", avsr_result.label,avsr_result.prob);//向txt中写入ASR识别结果
    fclose(asr_fp_result);

    //soc:send asr result to pc
    FILE *tmpAsrFp_result = fopen(asrCmdFinal_result, "rb"); // 以二进制方式打开（创建）文件
    while ((nCount_send = fread(send_buffer, 1, SEND_BUFFER_SIZE, tmpAsrFp_result)) > 0)
    {
        send(sock_cli, send_buffer, nCount_send, 0);
        // printf("===send vi_result len is %d==\n", nCount_send);
    }
    fclose(tmpAsrFp_result);
    printf("===soc end==\n");

    mtx_uart.lock();
    uart_main_func(1, 2, 0, 0);
    mtx_uart.unlock();
}


